package com.algrithom.heap;

import java.util.Arrays;
import java.util.Objects;
import java.util.PriorityQueue;

import com.common.constants.NumberConstants;

/**
 * 大顶堆和小顶堆的应用
 *
 * @author think
 * @version 1.0.0
 * @since 2020/2/6
 */
public class HeapApplication1 {
    
    /**
     * 利用小顶堆求最大的K个数
     *
     * @param nums 输入
     * @param k k个数
     * @return 符合条件的数据
     */
    public static int[] findKnumMax(int[] nums,int k){
        // 队列默认自然顺序排列，小顶堆，不必重写compare 默认是升序
        PriorityQueue<Integer> pq = new PriorityQueue<>(k);
        
        for (int num : nums) {
            if (pq.size() < k) {
                pq.offer(num);
                // 如果堆顶元素 < 新数，则删除堆顶，加入新数入堆
            } else if (Objects.nonNull(pq.peek()) && pq.peek() < num) {
                pq.poll();
                pq.offer(num);
            }
        }
        
        int[] result = new int[k];
        for (int i = 0; i < k && !pq.isEmpty(); i++) {
            result[i] = pq.poll();
        }
        return result;
    }
    
    /**
     * 利用大顶堆求最小的K个数
     *
     * @param nums 输入
     * @param k k个数
     * @return 符合条件的数据
     */
    public static int[] findKnumMin(int[] nums,int k){
        PriorityQueue<Integer> pq = new PriorityQueue<>(k,(o1,o2) -> o2 - o1);
        
        for (int num : nums) {
            if (pq.size() < k) {
                pq.offer(num);
                // 如果堆顶元素 > 新数，则删除堆顶，加入新数入堆
            } else if (Objects.nonNull(pq.peek()) && pq.peek() > num) {
                pq.poll();
                pq.offer(num);
            }
        }
        int[] result = new int[k];
        for (int i = 0; i < k && !pq.isEmpty(); i++) {
            result[i] = pq.poll();
        }
        return result;
    }
    
    /**
     * 求第k个最小元素/k个最大元素
     *
     * @param nums 输入
     * @param k 第K个
     * @return 第K个元素
     */
    public static int findKthLargest(int[] nums,int k){
        PriorityQueue<Integer> pq = new PriorityQueue<>(k,(o1,o2) -> o2 - o1);
        
        for (int val : nums) {
            pq.offer(val);
            // 维护堆的大小为 K
            if (pq.size() > k) {
                pq.poll();
            }
        }
        if (Objects.isNull(pq.peek())) {
            return 0;
        }
        return pq.peek();
    }
    
    private static void heapMax(){
        PriorityQueue<Integer> queue = new PriorityQueue<>(NumberConstants.TEN,(o1,o2) -> o2 - o1);
        Arrays.asList(1,10,6,7,9,8,2,4,3,5).forEach(queue::offer);
        System.out.println(Arrays.toString(queue.toArray()));
    }
    
    private static void heapMin(){
        PriorityQueue<Integer> queue = new PriorityQueue<>(NumberConstants.TEN);
        Arrays.asList(1,10,6,7,9,8,2,4,3,5).forEach(queue::offer);
        System.out.println(Arrays.toString(queue.toArray()));
    }
    
    public static void main(String[] args){
        int[] arr = new int[]{1,6,2,3,5,4,8,7,9};
        System.out.println(Arrays.toString(findKnumMax(arr,3)));
        System.out.println(Arrays.toString(findKnumMin(arr,3)));
        System.out.println(findKthLargest(arr,5));
        heapMax();
        heapMin();
    }
}
